"""
切比雪夫多项式KAN层可视化示例
演示ChebyKANLinear的工作原理和切比雪夫多项式的特性
"""

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from layers.ChebyKANLayer import ChebyKANLinear

def visualize_chebyshev_polynomials(max_degree=5):
    """
    可视化切比雪夫多项式
    
    Args:
        max_degree (int): 最高次数
    """
    x = np.linspace(-1, 1, 1000)
    
    plt.figure(figsize=(12, 8))
    
    for n in range(max_degree + 1):
        # 计算n次切比雪夫多项式
        T_n = np.cos(n * np.arccos(x))
        plt.plot(x, T_n, label=f'T_{n}(x)', linewidth=2)
    
    plt.xlabel('x')
    plt.ylabel('T_n(x)')
    plt.title('切比雪夫多项式 T_n(x) = cos(n·arccos(x))')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.axhline(y=0, color='k', linestyle='-', alpha=0.3)
    plt.axvline(x=0, color='k', linestyle='-', alpha=0.3)
    plt.ylim(-1.5, 1.5)
    plt.show()

def demonstrate_kan_layer():
    """
    演示ChebyKANLinear层的功能
    """
    # 创建KAN层
    input_dim = 2
    output_dim = 1
    degree = 4
    
    kan_layer = ChebyKANLinear(input_dim, output_dim, degree)
    
    # 生成测试数据
    x = torch.randn(100, input_dim)
    
    # 前向传播
    with torch.no_grad():
        output = kan_layer(x)
    
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"切比雪夫系数形状: {kan_layer.cheby_coeffs.shape}")
    print(f"多项式度数: {degree}")
    
    return kan_layer, x, output

def compare_activation_functions():
    """
    比较传统激活函数与切比雪夫多项式的表达能力
    """
    x = torch.linspace(-2, 2, 1000).unsqueeze(1)
    
    # 传统激活函数
    relu_output = torch.relu(x)
    tanh_output = torch.tanh(x)
    sigmoid_output = torch.sigmoid(x)
    
    # 切比雪夫KAN层
    kan_layer = ChebyKANLinear(1, 1, degree=5)
    with torch.no_grad():
        kan_output = kan_layer(x)
    
    plt.figure(figsize=(15, 10))
    
    # 绘制传统激活函数
    plt.subplot(2, 3, 1)
    plt.plot(x.numpy(), relu_output.numpy(), 'b-', linewidth=2)
    plt.title('ReLU激活函数')
    plt.grid(True, alpha=0.3)
    
    plt.subplot(2, 3, 2)
    plt.plot(x.numpy(), tanh_output.numpy(), 'g-', linewidth=2)
    plt.title('Tanh激活函数')
    plt.grid(True, alpha=0.3)
    
    plt.subplot(2, 3, 3)
    plt.plot(x.numpy(), sigmoid_output.numpy(), 'r-', linewidth=2)
    plt.title('Sigmoid激活函数')
    plt.grid(True, alpha=0.3)
    
    # 绘制切比雪夫KAN输出
    plt.subplot(2, 3, 4)
    plt.plot(x.numpy(), kan_output.numpy(), 'm-', linewidth=2)
    plt.title('切比雪夫KAN层输出')
    plt.grid(True, alpha=0.3)
    
    # 绘制切比雪夫多项式基函数
    plt.subplot(2, 3, 5)
    x_cheby = torch.linspace(-1, 1, 1000)
    for i in range(6):
        T_i = torch.cos(i * torch.acos(x_cheby))
        plt.plot(x_cheby.numpy(), T_i.numpy(), label=f'T_{i}', linewidth=2)
    plt.title('切比雪夫多项式基函数')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # 对比所有函数
    plt.subplot(2, 3, 6)
    plt.plot(x.numpy(), relu_output.numpy(), 'b-', label='ReLU', linewidth=2)
    plt.plot(x.numpy(), tanh_output.numpy(), 'g-', label='Tanh', linewidth=2)
    plt.plot(x.numpy(), sigmoid_output.numpy(), 'r-', label='Sigmoid', linewidth=2)
    plt.plot(x.numpy(), kan_output.numpy(), 'm-', label='KAN', linewidth=2)
    plt.title('激活函数对比')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()

def analyze_kan_coefficients():
    """
    分析KAN层的切比雪夫系数
    """
    # 创建并训练一个简单的KAN层
    kan_layer = ChebyKANLinear(3, 2, degree=4)
    
    # 模拟训练过程
    optimizer = torch.optim.Adam(kan_layer.parameters(), lr=0.01)
    
    # 生成训练数据
    x_train = torch.randn(1000, 3)
    y_train = torch.sin(x_train.sum(dim=1, keepdim=True)) + 0.1 * torch.randn(1000, 1)
    y_train = y_train.repeat(1, 2)  # 扩展到2维输出
    
    # 训练
    for epoch in range(100):
        optimizer.zero_grad()
        output = kan_layer(x_train)
        loss = nn.MSELoss()(output, y_train)
        loss.backward()
        optimizer.step()
        
        if epoch % 20 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.6f}")
    
    # 可视化系数
    coeffs = kan_layer.cheby_coeffs.detach().numpy()
    
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    
    for i in range(3):  # 输入维度
        for j in range(2):  # 输出维度
            ax = axes[j, i]
            ax.bar(range(len(coeffs[i, j])), coeffs[i, j])
            ax.set_title(f'输入{i+1} -> 输出{j+1}的切比雪夫系数')
            ax.set_xlabel('多项式次数')
            ax.set_ylabel('系数值')
            ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()
    
    return coeffs

def main():
    """
    主函数：运行所有演示
    """
    print("=== 切比雪夫多项式KAN层演示 ===\n")
    
    print("1. 可视化切比雪夫多项式")
    visualize_chebyshev_polynomials()
    
    print("\n2. 演示KAN层基本功能")
    kan_layer, x, output = demonstrate_kan_layer()
    
    print("\n3. 比较激活函数")
    compare_activation_functions()
    
    print("\n4. 分析KAN系数")
    coeffs = analyze_kan_coefficients()
    
    print("\n=== 演示完成 ===")

if __name__ == "__main__":
    main()